import torch
import torch.utils.data as Data
from TransFM.dataset import TestDataset, BookUserInfo
import json
from typing import Dict, List

from utils.log_output import simple_text_log


class Evaluator:
    def __init__(self, hyper_params, book_user_info: BookUserInfo):
        self.hyper_params = hyper_params
        
        self.dataset = TestDataset(hyper_params, book_user_info)

    def evaluate(self, model) -> List[List[int]]:
        print('evaluate: Evaluating...')
        result = [None] * self.hyper_params['user_cnt']

        loader = Data.DataLoader(self.dataset, self.hyper_params['eval_batch_size'], num_workers=8)
        model.eval()

        result_tensor = torch.zeros(len(self.dataset), dtype=torch.float)
        result_bookid = torch.zeros(len(self.dataset), dtype=torch.long)
        cur_index = 0

        with torch.no_grad():
            for step, data in enumerate(loader):
                user_id, school, grade, user_len, date, sem, pos_bid, pos_title, pos_type = [i.to(self.hyper_params['device']) for i in data]

                score = model(user_id, school, grade, user_len, date, sem, pos_bid, pos_title, pos_type)

                result_tensor[cur_index:cur_index+user_id.size(0)] = score.detach().cpu()
                result_bookid[cur_index:cur_index+user_id.size(0)] = pos_bid.detach().cpu()
                cur_index += user_id.size(0)

                if step % 100 == 0:
                    simple_text_log('train', f'Evaluation ({user_id[0].item()}/{self.dataset.user_cnt})...')

        result_tensor = result_tensor.reshape(self.hyper_params['user_cnt'], -1)
        result_bookid = result_bookid.reshape(self.hyper_params['user_cnt'], -1)

        result_tensor = torch.argsort(result_tensor, dim=-1, descending=True)
        result_bookid = torch.gather(result_bookid, dim=1, index=result_tensor)

        # Convert to result
        for i in range(self.hyper_params['user_cnt']):
            result[i] = result_bookid[i].tolist()

        model.train()
        return result
